import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
 
from scipy.signal import savgol_filter
 
from sklearn.cross_decomposition import PLSRegression
from sklearn.model_selection import KFold, cross_val_predict, train_test_split
from sklearn.metrics import accuracy_score

# Load data into a Pandas dataframe
data = pd.read_csv('./data/milk-powder.csv')
# Extract firsint and last label  a new dataframe
binary_data = data[(data['labels'] == 5 ) | (data['labels'] == 6)]
# Read data into a numpy array and apply simple smoothing
X_binary = savgol_filter(binary_data.values[:,2:], 15, polyorder = 3, deriv=0)
# Read categorical variables
y_binary = binary_data["labels"].values
# Map variables to 0 and 1
y_binary = (y_binary == 6).astype('uint8')

# Define the PLS regression object
pls_binary =PLSRegression(n_components=2)
# Fit and transform the data
X_pls = pls_binary.fit_transform(X_binary, y_binary)[0]

# Define the labels for the plot legend
labplot = ["60/40 ratio", "50/50 ratio"]
# Scatter plot
unique = list(set(y_binary))
colors = [plt.cm.jet(float(i)/max(unique)) for i in unique]
 
with plt.style.context(('ggplot')):
    plt.figure(figsize=(12,10))
    for i, u in enumerate(unique):
        col = np.expand_dims(np.array(colors[i]), axis=0)
        xi = [X_pls[j,0] for j in range(len(X_pls[:,0])) if y_binary[j] == u]
        yi = [X_pls[j,1] for j in range(len(X_pls[:,1])) if y_binary[j] == u]
        plt.scatter(xi, yi, c=col, s=100, edgecolors='k',label=str(u))
 
    plt.xlabel('Latent Variable 1')
    plt.ylabel('Latent Variable 2')
    plt.legend(labplot,loc='lower left')
    plt.title('PLS cross-decomposition')
    plt.show()
    
# Test-train split
X_train, X_test, y_train, y_test = train_test_split(X_binary, y_binary, test_size=0.2, random_state=19)
# Define the PLS object
pls_binary = PLSRegression(n_components=2)
# Fit the training set
pls_binary.fit(X_train, y_train)
 
# Predictions: these won't generally be integer numbers
y_pred = pls_binary.predict(X_test)[:,0]
# "Force" binary prediction by thresholding
binary_prediction = (pls_binary.predict(X_test)[:,0] > 0.5).astype('uint8')
print(binary_prediction, y_test)